//用于测试设备的IO性能，确定模型的性能参数
//Usage: cal_profiler -t 1 -m uint8 -d 128 -v 1000
//-t指线程数，-m指向量数据类型(float32 or uint8)，-d指向量维度，-v指每个线程计算的向量数量
#include <iostream>
#include <thread> 
#include <random>
#include <sys/stat.h>
#include "../annlite_utils.h"
#include "statis.h"

using namespace std; 

#define DISTANCE l2_distance_general //l2_distance_general包含多种实现，可以处理不同数据类型的计算

//向量随机生成器。如果是uint8_t类型则范围为0-255，float32类型则范围为0-1
//会直接返回内存地址，计算完毕后需释放
template <class T>
T* random_vector(int dim, int count, int data_type){ 
    std::random_device rd;
    std::mt19937 gen(rd());
    T* vector = new T[dim*count];
    if(data_type == 0){
        std::uniform_real_distribution<float> dis(0.0, 1.0);
        for(int i=0; i<count; i++){
            for(int j=0; j<dim; j++){
                vector[i*dim+j] = dis(gen);
            }
        }
    }
    else if(data_type == 1){
        std::uniform_int_distribution<int> dis(0, 255);
        for(int i=0; i<count; i++){
            for(int j=0; j<dim; j++){
                vector[i*dim+j] = dis(gen);
            }
        }
    }
    return vector;
}

template <class T>
float cal(T* query, T* vector, int dim, int count, Profiler_Statis& statis, int thread_id){ 
    //记录开始时间
    auto start_time = statis.get_time();
    float sum = 0.0;
    for(int i=0; i<count; i++){
        float distance = DISTANCE(query, vector+i*dim, dim);
        sum += distance;
    }
    //记录结束时间
    auto end_time = statis.get_time();
    double time = statis.elapsed_time(start_time, end_time);
    statis.add_time(thread_id, time);

    return sum;
}

int main(int argc, char* argv[])
{
    // 默认参数值
    int thread_count = 1; 
    long long total_vectors = 1000; //每个线程计算多少个向量
    int dim = 128; //向量的维度
    int data_type = 0; //0: float32, 1: uint8_t

    // 解析命令行参数
    for (int i = 1; i < argc; ++i) {
        string arg = argv[i];
        if (arg == "-v" && i + 1 < argc) {
            total_vectors = stoll(argv[++i]);
        } else if (arg == "-t" && i + 1 < argc) {
            thread_count = stoi(argv[++i]);
        } else if (arg == "-d" && i + 1 < argc) {
            dim = stoi(argv[++i]);
        } else if (arg == "-m" && i + 1 < argc) {
            string type = argv[++i];
            if(type == "uint8")
                data_type = 1;
            else if(type == "float32")
                data_type = 0;
            else{
                cout << "Error: data_type must be uint8 or float32" << endl;
                return 0;
            }
        } else if (arg == "-h") {
            cout << "Usage: " << argv[0] << " [-v total_vectors] [-t thread_count] [-d dim] [-m data_type]" << endl;
            return 0;
        }
    }

    cout << "Total vectors: " << total_vectors << endl;
    cout << "Dim: " << dim << endl;
    if(data_type == 0)
        cout << "Data type: float32" << endl;
    else if(data_type == 1)
        cout << "Data type: uint8_t" << endl;
    cout << "Thread num: " << thread_count << endl;
    Profiler_Statis statis(thread_count);

    if(data_type == 0){
        cout << "Start generate vector" << endl;
        auto start_time = statis.get_time();
        //生成查询向量
        auto query = random_vector<float>(dim, 1, data_type);
        //为每个线程生成计算的向量数据集
        vector<float*> vectors_list;
        for(int i=0; i<thread_count; i++){
            auto vectors = random_vector<float>(dim, total_vectors, data_type);
            vectors_list.push_back(vectors);
        }
        auto end_time = statis.get_time();
        double init_time = statis.elapsed_time(start_time, end_time);
        cout << "Init time: " << init_time << " ms" << endl;

        //开始测试
        cout << "Start IO performance profile" << endl;
        statis.start(); 
        
        std::vector<std::thread> threads;
        for (int i = 0; i < thread_count; i++) {
            // 修改捕获方式，按值捕获需要的变量
            threads.emplace_back([=, &statis]() {
                cal(query, vectors_list[i], dim, total_vectors, statis, i);
            });
            statis.add_io_count(i, total_vectors);
            statis.add_io_size(i, total_vectors*dim*sizeof(float));
        }

        for (auto& t : threads) {
            t.join();
        }

        statis.end();
        //输出性能数据
        statis.cal_print();

        free(query);
        for(int i = 0; i < thread_count; i++){
            free(vectors_list[i]);
        }
    }

    if(data_type == 1){
        cout << "Start generate vector" << endl;
        auto start_time = statis.get_time();
        //生成查询向量
        auto query = random_vector<uint8_t>(dim, 1, data_type);
        //为每个线程生成计算的向量数据集
        vector<uint8_t*> vectors_list;
        for(int i=0; i<thread_count; i++){
            auto vectors = random_vector<uint8_t>(dim, total_vectors, data_type);
            vectors_list.push_back(vectors);
        }
        auto end_time = statis.get_time();
        double init_time = statis.elapsed_time(start_time, end_time);
        cout << "Init time: " << init_time << " ms" << endl;

        //开始测试
        cout << "Start IO performance profile" << endl;
        statis.start(); 
        
        std::vector<std::thread> threads;
        for (int i = 0; i < thread_count; i++) {
            // 修改捕获方式，按值捕获需要的变量
            threads.emplace_back([=, &statis]() {
                cal(query, vectors_list[i], dim, total_vectors, statis, i);
            });
            statis.add_io_count(i, total_vectors);
            statis.add_io_size(i, total_vectors*dim*sizeof(float));
        }

        for (auto& t : threads) {
            t.join();
        }

        statis.end();
        //输出性能数据
        statis.cal_print();

        free(query);
        for(int i = 0; i < thread_count; i++){
            free(vectors_list[i]);
        }
    }

    return 0;
}